// main.go
package main

import (
	"context"
	"encoding/json"
	"fmt"
	admissionv1 "k8s.io/api/admission/v1"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
	"net/http"
	"os"
)

var (
	clientset *kubernetes.Clientset
)

func main() {
	config, err := rest.InClusterConfig()
	if err != nil {
		panic(err.Error())
	}
	clientset, err = kubernetes.NewForConfig(config)
	if err != nil {
		panic(err.Error())
	}

	http.HandleFunc("/mutate", mutateHandler)
	http.HandleFunc("/health", healthHandler)

	port := os.Getenv("PORT")
	if port == "" {
		port = "8443"
	}

	fmt.Printf("Server listening on :%s\n", port)
	err = http.ListenAndServeTLS(":"+port, "/etc/webhook/certs/tls.crt", "/etc/webhook/certs/tls.key", nil)
	if err != nil {
		panic(err)
	}
}

func healthHandler(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(http.StatusOK)
}

func mutateHandler(w http.ResponseWriter, r *http.Request) {
	admissionReview := &admissionv1.AdmissionReview{}
	if err := json.NewDecoder(r.Body).Decode(admissionReview); err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	pod := &corev1.Pod{}
	if err := json.Unmarshal(admissionReview.Request.Object.Raw, pod); err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	var patches []map[string]interface{}

	// 只有当Pod已经调度到节点时才处理
	if pod.Spec.NodeName != "" {
		node, err := clientset.CoreV1().Nodes().Get(context.Background(), pod.Spec.NodeName, metav1.GetOptions{})
		if err == nil {
			if volumePath, ok := node.Annotations["hostpath-volume"]; ok {
				// 添加volume
				volumePatch := map[string]interface{}{
					"op":   "add",
					"path": "/spec/volumes/-",
					"value": corev1.Volume{
						Name: "node-hostpath",
						VolumeSource: corev1.VolumeSource{
							HostPath: &corev1.HostPathVolumeSource{
								Path: volumePath,
							},
						},
					},
				}
				patches = append(patches, volumePatch)

				// 为所有容器添加volumeMount
				for i := range pod.Spec.Containers {
					mountPatch := map[string]interface{}{
						"op":   "add",
						"path": fmt.Sprintf("/spec/containers/%d/volumeMounts/-", i),
						"value": corev1.VolumeMount{
							Name:      "node-hostpath",
							MountPath: "/hostpath",
						},
					}
					patches = append(patches, mountPatch)
				}
			}
		}
	}

	patchBytes, _ := json.Marshal(patches)
	admissionResponse := &admissionv1.AdmissionResponse{
		UID:     admissionReview.Request.UID,
		Allowed: true,
		Patch:   patchBytes,
		PatchType: func() *admissionv1.PatchType {
			pt := admissionv1.PatchTypeJSONPatch
			return &pt
		}(),
	}

	respAdmissionReview := &admissionv1.AdmissionReview{
		TypeMeta: metav1.TypeMeta{
			APIVersion: "admission.k8s.io/v1",
			Kind:       "AdmissionReview",
		},
		Response: admissionResponse,
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(respAdmissionReview)
}
